import os
import asyncio
import time
from typing import Dict, Any, Optional

class TaskManager:
    """
    简单的任务管理器，用于处理异步任务
    当没有Celery时作为替代方案使用
    """
    def __init__(self):
        self.tasks = {}  # 任务状态存储: task_id -> task_info
    
    def generate_task_id(self):
        """生成唯一的任务ID"""
        return f"task_{int(time.time())}_{os.getpid()}_{len(self.tasks)}"
    
    def add_task(self, task_function, *args, **kwargs):
        """添加任务到管理器"""
        task_id = self.generate_task_id()
        task_info = {
            "id": task_id,
            "status": "PENDING",
            "result": None,
            "error": None,
            "start_time": time.time(),
            "function": task_function,
            "args": args,
            "kwargs": kwargs
        }
        
        self.tasks[task_id] = task_info
        
        # 创建异步任务执行
        asyncio.create_task(self._run_task(task_id, task_function, *args, **kwargs))
        
        return task_id
    
    async def _run_task(self, task_id, task_function, *args, **kwargs):
        """执行异步任务"""
        try:
            self.tasks[task_id]["status"] = "RUNNING"
            result = await task_function(*args, **kwargs)
            self.tasks[task_id]["status"] = "SUCCESS"
            self.tasks[task_id]["result"] = result
            self.tasks[task_id]["end_time"] = time.time()
        except Exception as e:
            self.tasks[task_id]["status"] = "FAILURE"
            self.tasks[task_id]["error"] = str(e)
            self.tasks[task_id]["end_time"] = time.time()
    
    def get_task_status(self, task_id):
        """获取任务状态"""
        if task_id not in self.tasks:
            return {"status": "UNKNOWN", "error": "任务不存在"}
        
        task_info = self.tasks[task_id]
        return {
            "id": task_info["id"],
            "status": task_info["status"],
            "result": task_info["result"] if task_info["status"] == "SUCCESS" else None,
            "error": task_info["error"] if task_info["status"] == "FAILURE" else None,
            "start_time": task_info["start_time"],
            "end_time": task_info.get("end_time")
        }
    
    def clean_old_tasks(self, max_age_hours=24):
        """清理旧任务"""
        current_time = time.time()
        max_age_seconds = max_age_hours * 3600
        
        task_ids = list(self.tasks.keys())
        for task_id in task_ids:
            task = self.tasks[task_id]
            # 如果任务已完成且超过最大保留时间，删除它
            if (task["status"] in ["SUCCESS", "FAILURE"] and 
                "end_time" in task and 
                current_time - task["end_time"] > max_age_seconds):
                del self.tasks[task_id]

# 全局任务管理器实例
task_manager = TaskManager()

try:
    from celery import Celery
    
    # 配置Celery
    celery_app = Celery('retrieval_tasks')
    celery_app.conf.broker_url = os.environ.get('CELERY_BROKER_URL', 'redis://localhost:6379/1')
    celery_app.conf.result_backend = os.environ.get('CELERY_RESULT_BACKEND', 'redis://localhost:6379/1')
    
    @celery_app.task
    def process_file_content_task(content, file_url, file_name):
        """处理文件内容的Celery任务"""
        import asyncio
        from kg_retrieval.optimized_embedding import OptimizedEmbedding
        from elasticsearch import AsyncElasticsearch
        
        # 初始化嵌入处理器
        optimized_embedder = OptimizedEmbedding()
        
        async def _process():
            # 处理文件内容
            file_data = await optimized_embedder.process_file_content(content, file_name, file_url)
            
            # 连接到ES并存储数据，这里使用环境变量或默认值
            es_hosts = os.environ.get('ELASTICSEARCH_HOSTS', 'https://localhost:9200').split(',')
            es_user = os.environ.get('ELASTICSEARCH_USERNAME', 'elastic')
            es_pass = os.environ.get('ELASTICSEARCH_PASSWORD', '')
            es_index = os.environ.get('ELASTICSEARCH_INDEX', 'jd_file_index_formal')
            
            try:
                es = AsyncElasticsearch(
                    hosts=es_hosts,
                    basic_auth=(es_user, es_pass),
                    verify_certs=False
                )
                
                # 存储到ES
                response = await es.index(index=es_index, document=file_data)
                await es.close()
                return {"status": "success", "doc_id": response["_id"]}
            except Exception as e:
                if 'es' in locals():
                    await es.close()
                return {"status": "error", "message": str(e)}
        
        # 运行异步函数
        loop = asyncio.get_event_loop()
        return loop.run_until_complete(_process())
    
    has_celery = True
except ImportError:
    has_celery = False
    print("Celery未安装，将使用内置任务管理器")

async def process_file_content(content, file_url, file_name):
    """处理文件内容的异步任务，兼容Celery和内置任务管理器"""
    from kg_retrieval.optimized_embedding import OptimizedEmbedding
    
    # 初始化嵌入处理器
    optimized_embedder = OptimizedEmbedding()
    
    # 处理文件内容
    try:
        file_data = await optimized_embedder.process_file_content(content, file_name, file_url)
        
        # 这里可以添加保存到本地文件或数据库的逻辑，作为Elasticsearch的替代
        # 简单起见，这里只返回处理结果
        return {
            "status": "success", 
            "message": "文件内容已处理",
            "file_name": file_name,
            "file_url": file_url
        }
    except Exception as e:
        return {"status": "error", "message": str(e)}

def submit_task(content, file_url, file_name):
    """提交一个处理文件的任务"""
    if has_celery:
        # 使用Celery
        task = process_file_content_task.delay(content, file_url, file_name)
        return {
            "task_id": task.id,
            "status": "PENDING",
            "message": "任务已提交到Celery"
        }
    else:
        # 使用内置任务管理器
        task_id = task_manager.add_task(process_file_content, content, file_url, file_name)
        return {
            "task_id": task_id,
            "status": "PENDING", 
            "message": "任务已提交到内置任务管理器"
        }

def get_task_status(task_id):
    """获取任务状态，兼容Celery和内置任务管理器"""
    if has_celery:
        task = celery_app.AsyncResult(task_id)
        if task.ready():
            if task.successful():
                return {
                    "status": "completed",
                    "result": task.result
                }
            else:
                return {
                    "status": "failed",
                    "error": str(task.result)
                }
        else:
            return {
                "status": "processing",
                "message": "任务正在处理中"
            }
    else:
        # 使用内置任务管理器
        return task_manager.get_task_status(task_id) 